import pickle as pkl
import numpy as np
import matplotlib.pyplot as plt
from pygenstability.pygenstability import _evaluate_VI

if __name__ == "__main__":
    data = pkl.load(open("hox_gene_expression_data.pkl", "rb"))[:-1]
    c = ["C0", "C1", "C2"]

    plt.figure(figsize=(5, 2))
    ax1 = plt.gca()

    for i, method in enumerate(["geometric_modularity", "markovstab"]):
        cluster_results = pkl.load(open(method + "_results.pkl", "rb"))
        times = cluster_results["times"]
        community_id = cluster_results["community_id"]

        ground_truth = np.unique(data["Neuron Class"].to_list(), return_inverse=True)[1]

        n_c = []
        vis = []
        for _i in range(len(times)):
            vis.append(_evaluate_VI((0, 1), [community_id[_i], ground_truth]))
            n_c.append(len(np.unique(community_id[_i])))

        ax1.plot(np.log10(times), vis, c=c[i], ls="-", label=method)
        plt.ylabel("VI")
        ax1.legend(loc="best")
        plt.axis([np.log10(times[0]), np.log10(times[-1]), 0, 0.85])
    plt.savefig("ground_truth_VI.svg")
    plt.show()
